import json
import os
import time

import javalang
import networkx as nx
from app.logic.utils.filepath import PDG_DIR, SRC_DIR

'''
按项目路径生成程序依赖图PDG
'''


class GeneratePDG:
    """
    生成程序依赖图 (PDG) 并使用 networkx 处理
    """

    def __init__(self, work_path_dir, project_path_dir):
        self.workspace = os.path.join(work_path_dir, str(project_path_dir))
        self.src_dir = os.path.join(self.workspace, SRC_DIR)
        self.pdg_dir = os.path.join(self.workspace, PDG_DIR)

        # 确保 PDG 存储目录存在
        if not os.path.exists(self.pdg_dir):
            os.makedirs(self.pdg_dir)

    def parse_java_file(self, java_file_path):
        """
        使用 javalang 解析 Java 文件，返回 AST
        """
        try:
            with open(java_file_path, 'r', encoding='utf-8') as f:
                code = f.read()
            tree = javalang.parse.parse(code)
            return tree
        except Exception as e:
            print(f"Failed to parse {java_file_path}: {e}")
            return None

    def build_pdg(self, tree):
        """
        根据 AST 构建程序依赖图 (PDG)
        """
        graph = nx.DiGraph()

        # 创建节点和边
        def add_node(node, parent=None):
            if isinstance(node, javalang.ast.Node):
                # 添加节点到图
                node_id = id(node)
                graph.add_node(node_id, label=type(node).__name__)

                # 添加父子关系边
                if parent is not None:
                    graph.add_edge(parent, node_id)

                # 递归处理子节点
                for child in node.children:
                    if isinstance(child, (list, tuple)):
                        for subchild in child:
                            add_node(subchild, node_id)
                    else:
                        add_node(child, node_id)
                return node_id
            elif isinstance(node, str):
                # 处理字符串节点
                node_id = id(node)
                graph.add_node(node_id, label=node)
                if parent is not None:
                    graph.add_edge(parent, node_id)
                return node_id

        # 遍历 AST 树并构建图
        if hasattr(tree, 'types'):
            for t in tree.types:
                add_node(t)

        return graph

    def save_pdg(self, graph, output_path):
        """
        保存图到文件
        """
        data = nx.node_link_data(graph)  # 转为 JSON 格式
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=4)
        print(f"PDG saved to {output_path}")

    def generate_pdg_for_file(self, java_file_path):
        """
        为单个 Java 文件生成 PDG
        """
        tree = self.parse_java_file(java_file_path)
        if tree is None:
            return

        pdg = self.build_pdg(tree)

        # 保存 PDG 图
        file_name = os.path.basename(java_file_path)
        pdg_file = os.path.join(self.pdg_dir, f"{file_name}.json")
        self.save_pdg(pdg, pdg_file)

    def generate_pdg_for_directory(self):
        """
        遍历目录中的所有 Java 文件并生成它们的 PDG
        """
        for root, _, files in os.walk(self.src_dir):
            for file in files:
                if file.endswith('.java'):
                    java_file = os.path.join(root, file)
                    self.generate_pdg_for_file(java_file)

    def extract_and_save_pdg(self):
        """
        打印开始时间，生成 PDG 并保存，最后打印结束时间和总耗时
        """
        start_time = time.time()
        print("PDG Extraction Start Time: " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 生成 PDG
        self.generate_pdg_for_directory()

        # 打印结束时间
        print("PDG extraction completed at " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 计算总耗时
        total_time = time.time() - start_time
        print("Total time: " + str(total_time // 3600) + "h" +
              str((total_time % 3600) // 60) + "m" +
              str((total_time % 3600) % 60) + "s")


def gen_pdg(work_dir, project_path):
    pdg_generator = GeneratePDG(work_dir, project_path)
    pdg_generator.extract_and_save_pdg()


# 使用方法示例
if __name__ == "__main__":
    work_dir = '../../../data'  # 请修改为实际路径
    project_path = '29'  # 请修改为实际项目路径

    pdg_generator = GeneratePDG(work_dir, project_path)
    pdg_generator.extract_and_save_pdg()
